import json
import logging
from configparser import ConfigParser
from typing import Optional, List, Dict, Any, AsyncGenerator
import asyncio
import time

# 导入用于问答的 maas 库
from maas.constants.http_method_name import HttpMethodName
from maas.http.api_client import ApiClient
from maas.model.api_request import ApiRequest

# --- 基础配置 ---
logging.basicConfig(format="%(asctime)s - %(levelname)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S",
                    level=logging.INFO)
chat_logger = logging.getLogger(__name__)

# --- 全局变量 ---
API_SECRET: Optional[str] = None
API_KEY: Optional[str] = None
CHAT_API_URL: Optional[str] = None
CHAT_MODEL_NAME: Optional[str] = None


# --- 问答API核心功能  ---
def call_llm_api(messages: List[Dict], stream: bool = True, max_tokens: int = 6000, temperature: float = 0.9) -> Any:
    if not all([API_SECRET, API_KEY, CHAT_API_URL, CHAT_MODEL_NAME]):
        chat_logger.error("聊天模型API配置(maas)未完全初始化")
        return None
    api_request = ApiRequest(HttpMethodName.POST, CHAT_API_URL)
    api_request.add_headers("Content-Type", "application/json")
    api_request.set_credential(API_KEY, API_SECRET)
    request_body = {
        "model": CHAT_MODEL_NAME,
        "messages": messages,
        "max_tokens": max_tokens,
        "temperature": temperature,
        "stream": stream
    }
    api_request.set_json_body(json.dumps(request_body))
    try:
        response = ApiClient().send_request(api_request)
        if not response.ok:
            chat_logger.error(f"maas API请求失败，状态码 {response.status_code}: {response.text}")
            return None
        return response if stream else response.json()
    except Exception as e:
        chat_logger.error(f"调用 maas API时发生异常: {e}", exc_info=True)
        return None


async def generate_chat_stream(messages: List[Dict]) -> AsyncGenerator[str, None]:
    response_obj = call_llm_api(messages, stream=True)
    if not response_obj:
        yield ""
        return
    try:
        buffer = ""
        for chunk in response_obj.iter_content(chunk_size=512, decode_unicode=True):
            buffer += chunk
            while '\n\n' in buffer:
                event_str, buffer = buffer.split('\n\n', 1)
                if not event_str.strip().startswith("data:"): continue
                json_str = event_str.strip()[5:].strip()
                if json_str == "[DONE]": break
                try:
                    data = json.loads(json_str)
                    content = data.get("choices", [{}])[0].get("delta", {}).get("content", "")
                    if content: yield content
                except (json.JSONDecodeError, IndexError):
                    pass
    except Exception as e:
        chat_logger.error(f"处理maas流式聊天响应时发生异常: {e}", exc_info=True)
    finally:
        if hasattr(response_obj, 'close'): response_obj.close()


# --- 长文本处理模块 ---

def split_text(text: str, chunk_size: int = 25000, overlap: int = 1000) -> List[str]:
    """将长文本按指定大小和重叠部分进行分片。"""
    if len(text) <= chunk_size: return [text]
    chunks = []
    start = 0
    while start < len(text):
        end = start + chunk_size
        chunks.append(text[start:end])
        if end >= len(text): break
        start = end - overlap
    chat_logger.info(f"文本被分割为 {len(chunks)} 个片段。")
    return chunks


def summarize_chunk_sync(chunk: str, chunk_number: int, total_chunks: int, final_instruction: str) -> str:
    """这是一个完全同步的函数，负责总结单个文本块。"""
    task_start_time = time.time()
    chat_logger.info(f"【线程启动】总结任务 {chunk_number}/{total_chunks}...")
    prompt = f"""你是一个智能的AI助手，正在协助处理一个大型文档。用户的最终目标是完成以下任务：
---
最终任务：{final_instruction}
---
你的当前工作是阅读下面这个文本片段（这是整个文档的第 {chunk_number}/{total_chunks} 部分），并提炼出所有与上述“最终任务”**最相关**的核心信息、关键论点和重要数据。你的总结必须服务于最终任务。
文本片段内容如下：
---
{chunk}
---
请针对“最终任务”输出你的总结："""
    messages = [{"role": "system", "content": "你是一个专注于提炼核心信息的文本总结助手。"},
                {"role": "user", "content": prompt}]
    response_obj = call_llm_api(messages, stream=True)
    if not response_obj: return f"第 {chunk_number} 部分总结失败。"
    full_summary = ""
    try:
        buffer = ""
        for content_chunk in response_obj.iter_content(chunk_size=512, decode_unicode=True):
            buffer += content_chunk
            while '\n\n' in buffer:
                event_str, buffer = buffer.split('\n\n', 1)
                if not event_str.strip().startswith("data:"): continue
                json_str = event_str.strip()[5:].strip()
                if json_str == "[DONE]": break
                try:
                    data = json.loads(json_str)
                    content = data.get("choices", [{}])[0].get("delta", {}).get("content", "")
                    if content: full_summary += content
                except (json.JSONDecodeError, IndexError):
                    pass
    finally:
        if hasattr(response_obj, 'close'): response_obj.close()
    task_end_time = time.time()
    duration = task_end_time - task_start_time
    chat_logger.info(f"【线程结束】总结任务 {chunk_number}/{total_chunks}，耗时: {duration:.2f} 秒。")
    return full_summary


# 关键改动：此函数现在是一个异步生成器
async def summarize_long_text_map_reduce(long_text: str, instruction: str) -> AsyncGenerator[str, None]:
    """
    使用 Map-Reduce 思想处理长文本，并将最终汇总结果以流式返回。
    """
    chunks = split_text(long_text)
    chat_logger.info(f"开始并行派发 {len(chunks)} 个总结任务到独立线程...")
    map_start_time = time.time()
    tasks = []
    for i, chunk in enumerate(chunks):
        task = asyncio.to_thread(
            summarize_chunk_sync, chunk, i + 1, len(chunks), instruction
        )
        tasks.append(task)
    chunk_summaries = await asyncio.gather(*tasks)
    map_end_time = time.time()
    total_map_duration = map_end_time - map_start_time
    chat_logger.info(f"所有 {len(chunks)} 个并行总结任务完成，总耗时: {total_map_duration:.2f} 秒。")

    combined_summaries = "\n\n---\n\n".join(
        f"这是第 {i + 1}/{len(chunks)} 部分的摘要（与最终任务相关）：\n{summary}" for i, summary in
        enumerate(chunk_summaries)
    )

    chat_logger.info("所有片段总结完成，正在进行最终汇总...")

    final_prompt = f"""你是一个顶级的AI分析师，负责将多个与特定任务相关的摘要，融合成一个逻辑连贯、信息全面的最终答案。
用户的最终任务是：
---
最终任务：{instruction}
---
下面是多个从原始文档中提炼出的、与此任务相关的摘要：
---
{combined_summaries}
---
请基于以上所有摘要，直接生成满足“最终任务”的完整、流畅的最终答案："""

    messages = [{"role": "system", "content": "你是一个高级AI分析师，擅长信息整合与解答。"},
                {"role": "user", "content": final_prompt}]

    # 关键改动：直接 yield generate_chat_stream 的结果
    # 这样，最终汇总的结果就会被实时地流式传输出去
    async for content_part in generate_chat_stream(messages):
        yield content_part

    chat_logger.info("长文本最终总结流式输出完成。")


# --- 业务流程编排  ---
async def process_simple_chat_stream(text: str, instruction: str, history: List[Dict]) -> AsyncGenerator[str, None]:
    """
    主聊天处理流程，升级了长文本处理逻辑，支持最终结果的流式输出。
    """
    MAX_LENGTH = 25000

    try:
        # 关键改动：将主逻辑包裹在 try/except 中
        if len(text) > MAX_LENGTH:
            chat_logger.info(f"检测到长文本 (长度: {len(text)} > {MAX_LENGTH})，启动目标导向的并行总结流程...")
            # 直接遍历长文本处理生成器，并将结果实时返回
            async for chunk in summarize_long_text_map_reduce(text, instruction):
                yield json.dumps({"code": 200, "msg": "生成中...", "data": {"text": chunk}}, ensure_ascii=False)
        else:
            # 对于短文本，流程保持不变
            messages = build_prompt_with_context(instruction, history, text)
            async for chunk in generate_chat_stream(messages):
                yield json.dumps({"code": 200, "msg": "生成中...", "data": {"text": chunk}}, ensure_ascii=False)

        # 统一在最后发送完成信号
        yield json.dumps({"code": 200, "msg": "完成", "data": {"text": ""}}, ensure_ascii=False)

    except Exception as e:
        chat_logger.error(f"处理聊天流时发生异常: {e}", exc_info=True)
        yield json.dumps({"code": 500, "msg": f"生成错误: {str(e)}", "data": None}, ensure_ascii=False)


# build_prompt_with_context 不再被长文本流程直接使用，但为短文本流程保留
def build_prompt_with_context(query: str, history: List[Dict], context_text: Optional[str]) -> List[Dict]:
    system_prompt = "你是一个知网的AI智能问答助手，你的任务是优先根据提供的参考信息来回答问题。"
    messages = [{"role": "system", "content": system_prompt}]
    if context_text:
        messages.append({"role": "system", "content": f"参考信息如下：\n\n---\n\n{context_text}\n\n---"})
    messages.extend(history)
    final_user_instruction = f"""
请优先根据我提供的“参考信息”来回答用户的问题。当文本中包含问题的回复时，你的回答必须精确、完整地基于所给的文本内容，当不包含时，请使用你的知识库进行回复。

**信息整合与输出规则：**
1.  **优先依赖参考信息**：你的核心任务是理解、梳理并整合“参考信息”中的内容，即使信息是多个片段，也需要将其组织成连贯、通顺且逻辑清晰的回答。
2.  **知识库使用条件**：仅在以下两种情况同时满足时，直接使用你自己的知识库进行回复：
    *   “参考信息”中未提及问题答案。
    *   “参考信息”内容为空。
    根据你自己的知识直接回复问题的答案，不要输出说明未提供参考信息等内容。

**参考信息中没有相关回复的特殊规则：**
1.  **根据你知识库的内容进行回复时**：不要提及‘参考信息’的任何信息，不要输出“参考信息中不存在相关内容”、“我将根据我自己的知识库进行回复。”、“您的问题与参考信息无关”等内容。根据你的知识回复问题答案。
2.  **仅当你自己也不能回复该问题时**：不要提及‘参考信息’的任何信息，响应‘我的知识库中没有相关内容，请提供更多参考信息。'

请开始处理用户的问题。
问题：{query}
"""
    messages.append({"role": "user", "content": final_user_instruction})
    return messages


# --- 统一的初始化函数  ---
def init_models(config: ConfigParser):
    global API_SECRET, API_KEY, CHAT_API_URL, CHAT_MODEL_NAME
    try:
        chat_config = config['chat_model']
        API_SECRET = chat_config.get('api_secret')
        API_KEY = chat_config.get('api_key')
        CHAT_API_URL = chat_config.get('api_url')
        CHAT_MODEL_NAME = chat_config.get('model_name')
        if not all([API_SECRET, API_KEY, CHAT_API_URL, CHAT_MODEL_NAME]):
            raise ValueError("chat_model配置不完整")
        chat_logger.info(f"maas聊天模型配置加载成功。模型: {CHAT_MODEL_NAME}, URL: {CHAT_API_URL}")
    except Exception as e:
        chat_logger.error(f"从 config.ini 初始化maas聊天模型失败: {e}")
